from django.http import Http404
from rest_framework import exceptions
from rest_framework.exceptions import PermissionDenied
from rest_framework.response import Response
from rest_framework.views import set_rollback


class ReturnMsg:
    def __init__(self, code=1, msg="成功", errors=None, data=None):
        self.code = code
        self.msg = msg
        self.errors = "" if errors is None else errors
        self.data = [] if data is None else data

    def dict(self):
        return {
            'code': self.code,
            'msg': self.msg,
            'errors': self.errors,
            'data': self.data
        }


def exception_handler(exc, context):
    if isinstance(exc, Http404):
        exc = exceptions.NotFound()
    elif isinstance(exc, PermissionDenied):
        exc = exceptions.PermissionDenied()

    if isinstance(exc, exceptions.APIException):
        headers = {}
        if getattr(exc, 'auth_header', None):
            headers['WWW-Authenticate'] = exc.auth_header
        if getattr(exc, 'wait', None):
            headers['Retry-After'] = '%d' % exc.wait

        if isinstance(exc.detail, (list, dict)):
            if isinstance(exc.detail, list):
                error = exc.detail[0]
            else:
                for k, v in exc.detail.items():
                    error = f"{k}: {v[0]}"
                    break
        else:
            error = exc.detail

        set_rollback()
        return Response({'code': 0, 'msg': "失败", 'error': error, 'data': []}, status=exc.status_code, headers=headers)

    return None
